# -*- coding: utf-8 -*-
import numpy as np
import os
import wrf_output
import gauge
import matplotlib.pyplot as plt


def Decumulation(seq):
    result = [seq[i + 1] - seq[i] for i in range(len(seq) - 1)]
    result.insert(0, seq[0])
    return result


def enlarge_xy_label():
    for x_label in plt.gca().xaxis.get_ticklabels():
        x_label.set_fontsize(20)
    for y_label in plt.gca().yaxis.get_ticklabels():
        y_label.set_fontsize(20)


def get_statistic(path, domain, startTime, endTime):
    # domains = ['01', '02']
    # get every mp_cu group wrf_output and extract the center value
    mp = ['2', '6', '7']
    cu = ['0', '1', '2', '10']
    split_path = path.split('/')
    gridSpacing = split_path[-2]
    gauge_loc = ''
    if gridSpacing == '139':
        gauge_loc = r'F:/Research/Data/Gauge/Gauge_row_col_1.csv'
    elif gridSpacing == '51545':
        gauge_loc = r'F:/Research/Data/Gauge/Gauge_row_col_5.csv'
    elif gridSpacing == '103090':
        gauge_loc = r'F:/Research/Data/Gauge/Gauge_row_col_10.csv'
    results = []
    for i in mp:
        for j in cu:
            mp_cu = i + j
            if os.path.exists(path + mp_cu + '/'):
                input = path + mp_cu + '/wrfout_d' + domain + '_' + startTime
                all_gauges = wrf_output.GetAll(input, gauge_loc)
                all_gauges = Decumulation(all_gauges)
                results.append(all_gauges)
    results = np.asarray(results)
    min = np.min(results, axis=0)
    max = np.max(results, axis=0)
    band = np.array([min, max])
    band_50 = np.mean(band, axis=0)
    band_50 = band_50.astype(np.float64)
    # get the center gauge value
    gauge50 = gauge.Get50Gauge(startTime, endTime)

    # calculate dispersion statistic
    band_95 = np.sum(band, axis=0) * 0.95
    band_95 = band_95.astype(np.float64)
    band_5 = np.sum(band, axis=0) * 0.05
    band_5 = band_5.astype(np.float64)
    # dispersion statistic
    dispersion_d = np.mean((band_95 - band_5), axis=0)
    # normalized dispersion statistic
    # denominator = band_50.copy()
    # denominator[denominator == 0] = np.nan
    # nd = (band_95 - band_5) / denominator
    # masked_nd = np.ma.masked_array(nd, np.isnan(nd))
    # dispersion_nd = np.mean(masked_nd, axis=0).filled(np.nan)
    # print result
    # print("Dispersion statistic:\n", dispersion_d)
    # print("Normalized dispersion statistic:\n", dispersion_nd)
    # calculate MAE and MSE
    mae = np.abs(np.mean((band_50 - gauge50.values), axis=0))
    mse = np.mean(np.power((band_50 - gauge50.values), 2), axis=0)
    # print("MAE:\n", mae)
    # print("MSE:\n", mse)
    # output statistic result
    all_statistic = np.array([dispersion_d, mae, mse])
    all_statistic = all_statistic.T
    return all_statistic
    # output_path = r'F:/Experiment/Research/Result/' + \
    #     split_path[3] + '/' + split_path[4] + '/' + \
    #     split_path[5] + '/' + split_path[6] + '/'
    # if not os.path.exists(output_path):
    #     os.makedirs(output_path)
    # np.savetxt(output_path + split_path[6] + '.csv', all_statistic,
    #            delimiter=',', fmt='%.2f', header='Dispersion,MAE,MSE')


def main(path, domain, startTime, endTime):
    scenarios = ['139', '51545', '103090']
    scenarios_stat = []
    dispersion_median = []
    mae_median = []
    for scenario in scenarios:
        input_path = path + scenario + '/'
        stat = get_statistic(input_path, domain, startTime, endTime)
        dispersion_median.append(np.median(stat[:, 0]))
        mae_median.append(np.median(stat[:, 1]))
        scenarios_stat.append(stat)
    scenarios_stat = np.array(scenarios_stat)
    scenario_num, gauge_num, stat_num = scenarios_stat.shape

    all_statistic = np.zeros((scenario_num * gauge_num, stat_num))
    for i in range(gauge_num):
        gauge_stat = scenarios_stat[:, i, :]
        # convert the matrix of (scenario_num, gauge_num, stat_num) to
        # the matrix of (scenario_num * gauge_num, stat_num)
        if i == 0:
            row = i
        else:
            row = row + scenario_num
        for j in range(scenario_num):
            all_statistic[row + j] = gauge_stat[j]
    # output and save as csv
    # split_path = path.split('/')
    # output_path = r'F:/Research/ProcessResult/' + \
    #     split_path[3] + '/' + split_path[4] + '/' + \
    #     split_path[5] + '/csv/'
    # if not os.path.exists(output_path):
    #     os.makedirs(output_path)
    # np.savetxt(output_path + 'statistic.csv', all_statistic,
    #            delimiter=',', fmt='%.2f', header='Dispersion,MAE,MSE')

    # plot dispersion, mae, mse
    # plt.figure(figsize=(10, 10))
    # for k in range(0, all_statistic.shape[0], 3):
    #     # plot dispersion
    #     plt.subplot(2, 2, 1)
    #     plt.xlim(0.5, 3.5)
    #     plt.plot((1, 2, 3), (all_statistic[k, 0],
    #                          all_statistic[k + 1, 0],
    #                          all_statistic[k + 2, 0]), 'bo-')
    #     if k == 0:
    #         plt.xticks((1, 2, 3))
    #         plt.xlabel('Scenarios', fontdict={'size': 22})
    #         plt.ylabel('Dispersion', fontdict={'size': 22})
    #         enlarge_xy_label()
    #         # plot median
    #         plt.subplot(2, 2, 2)
    #         plt.xlim(0.5, 3.5)
    #         plt.plot((1, 2, 3), dispersion_median, 'ro-')
    #         plt.xticks((1, 2, 3))
    #         plt.xlabel('Scenarios', fontdict={'size': 22})
    #         plt.ylabel('median of Dispersion', fontdict={'size': 22})
    #         enlarge_xy_label()
    #     # plot bias
    #     plt.subplot(2, 2, 3)
    #     plt.xlim(0.5, 3.5)
    #     plt.plot((1, 2, 3), (all_statistic[k, 1],
    #                          all_statistic[k + 1, 1],
    #                          all_statistic[k + 2, 1]), 'go-')
    #     if k == 0:
    #         plt.xticks((1, 2, 3))
    #         plt.xlabel('Scenarios', fontdict={'size': 22})
    #         plt.ylabel('MAE', fontdict={'size': 22})
    #         enlarge_xy_label()
    #         # plot median
    #         plt.subplot(2, 2, 4)
    #         plt.xlim(0.5, 3.5)
    #         plt.plot((1, 2, 3), mae_median, 'ro-')
    #         plt.xticks((1, 2, 3))
    #         plt.xlabel('Scenarios', fontdict={'size': 22})
    #         plt.ylabel('median of MAE', fontdict={'size': 22})
    #         enlarge_xy_label()
    # plt.show()


if __name__ == "__main__":
    # new run
    # path = r'F:/Research/WRF_Output/2008/01/1912/'
    # domain = r'03'
    # startTime = r'2008-01-19_12_00_00'
    # endTime = r'2008-01-22_00_00_00'
    # main(path, domain, startTime, endTime)
    path = r'F:/Research/WRF_Output/2008/03/1500/'
    domain = r'03'
    startTime = r'2008-03-15_00_00_00'
    endTime = r'2008-03-16_12_00_00'
    main(path, domain, startTime, endTime)
    # path = r'F:/Research/WRF_Output/2008/09/2900/'
    # domain = r'03'
    # startTime = r'2008-09-29_00_00_00'
    # endTime = r'2008-10-02_06_00_00'
    # main(path, domain, startTime, endTime)
    # path = r'F:/Research/WRF_Output/2008/12/1200/'
    # domain = r'03'
    # startTime = r'2008-12-12_00_00_00'
    # endTime = r'2008-12-14_06_00_00'
    # main(path, domain, startTime, endTime)
